
import numpy as np
import matplotlib.pyplot as plt


def harris_corner(img):
    def trans_to_gray(img):
        gray = 0.2126 * img[..., 2] + 0.7152 * img[..., 1] + 0.0722 * img[..., 0]
        gray = gray.astype(np.uint8)
        return gray

    def sobel_filter(gray):
        height, weight = gray.shape

        sobel_y = np.array(((3, 10, 3), (0, 0, 0), (-3, -10, -3)), np.float32)
        sobel_x = np.array(((3, 0, -3), (10, 0, -10), (3, 0, -3)), np.float32)

        tmp = np.pad(gray, (1, 1), 'edge')

        Ix = np.zeros_like(gray, np.float32)
        Iy = np.zeros_like(gray, np.float32)

        for y in range(height):
            for x in range(weight):
                Ix[y, x] = np.mean(tmp[y: y + 3, x: x + 3] * sobel_x)
                Iy[y, x] = np.mean(tmp[y: y + 3, x: x + 3] * sobel_y)

        Ix2 = Ix ** 2
        Iy2 = Iy ** 2
        Ixy = Ix * Iy

        return Ix2, Iy2, Ixy

    def gaussian_filter(I, k=3, sigma=5):
        H, W = I.shape

        I_t = np.pad(I, (k // 2, k // 2), 'edge')

        # gaussian kernel
        K = np.zeros((k, k), np.float)
        for x in range(k):
            for y in range(k):
                _x = x - k // 2
                _y = y - k // 2
                K[y, x] = np.exp(-(_x ** 2 + _y ** 2) / (2 * (sigma ** 2)))
        K /= (sigma * np.sqrt(2 * np.pi))
        K /= K.sum()

        # filtering
        for y in range(H):
            for x in range(W):
                I[y, x] = np.sum(I_t[y: y + k, x: x + k] * K)

        return I

        # corner detect

    def corner_detect(gray, Ix2, Iy2, Ixy, k=0.04, th=0.1):
        res_list = []
        outy = np.array((gray, gray, gray))
        outy = np.transpose(outy, (1, 2, 0))

        R = (Ix2 * Iy2 - Ixy ** 2) - k * ((Ix2 + Iy2) ** 2)
        print(R)
        print(R[3, 6])
        print(R >= np.max(R) * th)
        outy[R >= np.max(R) * th] = [0, 0, 255]
        y, x = R.shape
        for i in range(1, y - 1):
            for j in range(1, x - 1):

                if R[i, j] >= np.max(R) * th and R[i, j] >= np.max(R[i - 1:i + 2, j - 1:j + 2]):
                    res_list.append((i, j))

        print(res_list)

        outy = outy.astype(np.uint8)

        return outy, res_list

        # 1. grayscale

    gray = trans_to_gray(img)

    Ix2, Iy2, Ixy = sobel_filter(gray)

    Ix2 = gaussian_filter(Ix2, 3, 3)
    Iy2 = gaussian_filter(Iy2, 3, 3)
    Ixy = gaussian_filter(Ixy, 3, 3)

    out, list1 = corner_detect(gray, Ix2, Iy2, Ixy)
    print(out)
    return out, list1

    # Read image
